#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>

// 定义常量
#define M 1604        // A 的行数
#define N 50          // A 的列数，C 的行数和列数
#define NUM_MATRICES 1536  // 512 * 3
#define TILE_SIZE 16  // 共享内存分块大小

// CUDA 核函数：计算 C = A^T * A
__global__ void matMulShared(const float *A_all, float *C_all, int m, int n, int num_matrices) {
    int matrix_id = blockIdx.z;  // 矩阵索引
    if (matrix_id >= num_matrices) return;

    int row = blockIdx.y * TILE_SIZE + threadIdx.y;  // C 的行索引
    int col = blockIdx.x * TILE_SIZE + threadIdx.x;  // C 的列索引

    const float *A = A_all + matrix_id * m * n;  // 当前矩阵 A 的起始地址
    float *C = C_all + matrix_id * n * n;        // 当前矩阵 C 的起始地址

    __shared__ float As[TILE_SIZE][TILE_SIZE];  // 存储 A^T 的子块（A 的列块）
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];  // 存储 A 的子块（行块）

    float Csub = 0.0f;

    // 分块计算
    for (int t = 0; t < (m + TILE_SIZE - 1) / TILE_SIZE; t++) {
        int k_start = t * TILE_SIZE;

        // 加载 A^T 的子块（A 的列块）
        if (row < n && k_start + threadIdx.x < m) {
            As[threadIdx.y][threadIdx.x] = A[(k_start + threadIdx.x) * n + row];
        } else {
            As[threadIdx.y][threadIdx.x] = 0.0f;
        }

        // 加载 A 的子块（行块）
        if (col < n && k_start + threadIdx.y < m) {
            Bs[threadIdx.y][threadIdx.x] = A[(k_start + threadIdx.y) * n + col];
        } else {
            Bs[threadIdx.y][threadIdx.x] = 0.0f;
        }

        __syncthreads();

        // 计算当前子块的乘积并累加
        for (int k = 0; k < TILE_SIZE && k_start + k < m; k++) {
            Csub += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        }

        __syncthreads();
    }

    // 写入结果
    if (row < n && col < n) {
        C[row * n + col] = Csub;
    }
}

int main() {
    const int m = M;  // 1604
    const int n = N;  // 50
    const int num_matrices = NUM_MATRICES;  // 1536
    const size_t A_size = num_matrices * m * n * sizeof(float);
    const size_t C_size = num_matrices * n * n * sizeof(float);

    // 分配主机内存
    float *h_A_all = (float *)malloc(A_size);
    float *h_C_all = (float *)malloc(C_size);

    // 初始化输入矩阵 A（示例数据）
    for (int i = 0; i < num_matrices * m * n; i++) {
        h_A_all[i] = (float)(i % 100) / 100.0f;
    }

    // 分配设备内存
    float *d_A_all, *d_C_all;
    cudaMalloc(&d_A_all, A_size);
    cudaMalloc(&d_C_all, C_size);

    // 设置线程块和网格
    dim3 blockDim(TILE_SIZE, TILE_SIZE);
    dim3 gridDim((n + TILE_SIZE - 1) / TILE_SIZE, (n + TILE_SIZE - 1) / TILE_SIZE, num_matrices);

    // 创建 CUDA 事件用于计时
    cudaEvent_t start, stop, copy_start, copy_stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventCreate(&copy_start);
    cudaEventCreate(&copy_stop);

    // 统计变量
    float total_time = 0.0f, total_copy_time = 0.0f, total_kernel_time = 0.0f;
    float max_time = 0.0f, max_copy_time = 0.0f, max_kernel_time = 0.0f;
    float min_time = 1e9f, min_copy_time = 1e9f, min_kernel_time = 1e9f;

    // 运行 1000 次
    for (int run = 0; run < 1000; run++) {
        // 测量拷贝时间
        cudaEventRecord(copy_start);
        cudaMemcpy(d_A_all, h_A_all, A_size, cudaMemcpyHostToDevice);
        cudaEventRecord(copy_stop);
        cudaEventSynchronize(copy_stop);
        float copy_time;
        cudaEventElapsedTime(&copy_time, copy_start, copy_stop);

        // 测量 CUDA 乘法时间
        cudaEventRecord(start);
        matMulShared<<<gridDim, blockDim>>>(d_A_all, d_C_all, m, n, num_matrices);
        cudaEventRecord(stop);
        cudaEventSynchronize(stop);
        float kernel_time;
        cudaEventElapsedTime(&kernel_time, start, stop);

        // 计算总时间
        float elapsed_time = copy_time + kernel_time;

        // 输出每次运行的时间
        printf("Run %d: 总时间=%.6f ms, 拷贝时间=%.6f ms, CUDA乘法时间=%.6f ms\n",
               run, elapsed_time, copy_time, kernel_time);

        // 更新统计数据
        total_time += elapsed_time;
        total_copy_time += copy_time;
        total_kernel_time += kernel_time;
        max_time = (elapsed_time > max_time) ? elapsed_time : max_time;
        max_copy_time = (copy_time > max_copy_time) ? copy_time : max_copy_time;
        max_kernel_time = (kernel_time > max_kernel_time) ? kernel_time : max_kernel_time;
        min_time = (elapsed_time < min_time) ? elapsed_time : min_time;
        min_copy_time = (copy_time < min_copy_time) ? copy_time : min_copy_time;
        min_kernel_time = (kernel_time < min_kernel_time) ? kernel_time : min_kernel_time;
    }

    // 输出统计结果
    printf("\n统计结果（1000次运行）：\n");
    printf("平均总时间: %.6f ms\n", total_time / 1000.0f);
    printf("平均拷贝时间: %.6f ms\n", total_copy_time / 1000.0f);
    printf("平均CUDA乘法时间: %.6f ms\n", total_kernel_time / 1000.0f);
    printf("最大总时间: %.6f ms\n", max_time);
    printf("最大拷贝时间: %.6f ms\n", max_copy_time);
    printf("最大CUDA乘法时间: %.6f ms\n", max_kernel_time);
    printf("最小总时间: %.6f ms\n", min_time);
    printf("最小拷贝时间: %.6f ms\n", min_copy_time);
    printf("最小CUDA乘法时间: %.6f ms\n", min_kernel_time);

    // 释放资源
    cudaEventDestroy(start);
    cudaEventDestroy(stop);
    cudaEventDestroy(copy_start);
    cudaEventDestroy(copy_stop);
    cudaFree(d_A_all);
    cudaFree(d_C_all);
    free(h_A_all);
    free(h_C_all);

    return 0;
}